import os
from PIL import Image
import numpy as np
from logger import setup_logger
import cv2
import time
logger = setup_logger()
def get_depth_image(image_path):
    # 提取文件名
    img_name = os.path.basename(image_path)

    # 获取 Color 上层目录，然后构建对应的 Depth 路径
    base_dir = os.path.dirname(os.path.dirname(image_path))
    
    # 替换路径中的 Color -> Depth 或 Color1 -> Depth1
    sub_dir = os.path.basename(os.path.dirname(image_path))  # 比如 Color 或 Color1
    if "Color" in sub_dir:
        depth_sub_dir = sub_dir.replace("Color", "Depth")
    else:
        print(f"未识别的子目录: {sub_dir}")
        return None

    depth_path = os.path.join(base_dir, depth_sub_dir, img_name)

    # 检查文件是否存在
    if not os.path.exists(depth_path):
        # 尝试将 .jpg 替换为 .png
        if img_name.endswith('.jpg'):
            png_name = img_name.replace('.jpg', '.png')
            depth_path = os.path.join(base_dir, depth_sub_dir, png_name)
            if not os.path.exists(depth_path):
                print(f"深度图路径不存在: {depth_path}")
                return None

    try:
        # 使用 PIL 打开图像
        pil_img = Image.open(depth_path)

        # 处理不同格式的深度图
        if pil_img.mode == 'I;16':  # 16位小端深度图
            depth_array = np.array(pil_img, dtype=np.uint16)
            # 可选：检测是否需要字节序转换（一般不需要）
            if np.max(depth_array) > 65535:
                depth_array = depth_array.byteswap().newbyteorder()
        elif pil_img.mode == 'I':  # 32位整数深度图
            depth_array = np.array(pil_img, dtype=np.uint32)
            depth_array = (depth_array >> 16).astype(np.uint16)  # 提取高16位
        else:
            depth_array = np.array(pil_img)
            print(f"警告: 未知模式 {pil_img.mode}")

        # 返回单通道图
        if depth_array.ndim == 2:
            return depth_array
        else:
            print(f"警告: 意外通道数 {depth_array.shape}")
            return depth_array[:, :, 0]

    except Exception as e:
        print(f"深度图处理失败: {str(e)}")
        return None
    
import os
import pickle
import cv2
from logger import setup_logger

logger = setup_logger()

def setup_standard_position(seg_pose, compute_position):
    """设置标准位置 - 支持缓存结果到文件"""
    # 确保缓存目录存在
    cache_dir = "./data/cache"
    os.makedirs(cache_dir, exist_ok=True)
    
    # 缓存文件路径
    seg_cache_path = os.path.join(cache_dir, "st_seg_result.pkl")
    pose_cache_path = os.path.join(cache_dir, "st_pose_result.pkl")
    
    # 尝试从缓存加载结果
    if os.path.exists(seg_cache_path) and os.path.exists(pose_cache_path):
        try:
            with open(seg_cache_path, "rb") as f:
                seg_result = pickle.load(f)
            with open(pose_cache_path, "rb") as f:
                pose_result = pickle.load(f)
            logger.info("==========从缓存加载标准位置结果成功==========")
        except Exception as e:
            logger.warning(f"缓存加载失败，重新计算: {e}")
            seg_result, pose_result = None, None
    else:
        seg_result, pose_result = None, None
    
    # 需要重新计算的情况
    logger.info("==========开始计算标准位置==========")
    st_img = cv2.imread("./data/Color/1.png")
    st_depth_img = get_depth_image("./data/Color/1.png")
    if seg_result is None or pose_result is None:
        if seg_pose is None:
            logger.error("==========请提供seg_pose实例==========")
            return None, None
        # 计算分割和姿态结果
        seg_pose.put(st_img)
        time.sleep(1)
        result = seg_pose.get_result(st_img)
        seg_result = result['seg_result'][0][0]
        pose_result = result['pose_result'][0] 
        if seg_result is None or pose_result is None:
            logger.error("==========未检测到标准图片的seg或pose结果==========")
            return None, None
        
        # 保存结果到缓存
        try:
            with open(seg_cache_path, "wb") as f:
                pickle.dump(seg_result, f)
            with open(pose_cache_path, "wb") as f:
                pickle.dump(pose_result, f)
            logger.info("==========标准位置结果已保存到缓存==========")
        except Exception as e:
            logger.error(f"缓存保存失败: {e}")
    
    # 设置标准位置
    # pose={"xy":pose_result.xy,"conf":pose_result.conf}
    # print
    st_pose, st_vector = compute_position.set_standard(
        st_img, st_depth_img, seg_result, pose_result
    )
    if st_pose is None or st_vector is None:
        logger.error("==========标准位置设定失败==========")
        return None, None
    
    logger.info("==========标准位置设定成功==========")
    logger.info(f"标准位置: {st_pose}")
    logger.info(f"标准向量: {st_vector}")
    return st_pose, st_vector
    
